# -*- coding:utf-8 -*-


# IMPORTANT: nodes' connection information is maintained by both of node class and the graph class
# The difference of those is the former is for the message-passing to the adjacent nodes( different types)\
# while the latter is for the topology and high-level thinking.


from  builtins import range
from functools import reduce
import numpy as np
from cmsg import *
import logging

class cNode(object):
    '''
        base class of node for a undirectional graph
        It contains the following data structure:
        1. basic information of this node:
            name: name string. it doesn't need to be exclusive for a graph.
            type: type of the node(string), factnode or varnode

        2. connectiveness information 
            adjNodes: List of the adjacent node instances of this node.
            adjNodeConnectionTypes: List of the adjacent node type, 'in','out' or 'bi' 

        3. message-passiing-relavent data
            msgClass: class of the output messages 
            inMsgs: Dict-type input messages, key is the index of the adjacent node in adjNodes
            outMsgs: Dict-type output messages. 
            enabled ： identify whether it can transport information

        Process of message passing
    '''

    def __init__(self,name):
        self.name =  name
        self.inMsgs = dict({})
        self.outMsgs = dict({})

        self.adjNodes = list([])
        self.adjNodeConnectionTypes =list([])
        self.adjSameNodes =list([])  # the adjacent same nodes
        self.type = ''
        self.enable = True


    def getAdjNodeIndex(self,adjNode):
        '''
            get adjacent nodes index in self.adjNodes
            adjNode is a string of node.name or instance of node
        '''
        if isinstance(adjNode,cNode):
            return self.adjNodes.index(adjNode)
        elif isinstance(adjNode,str):
            print 'this is not ready'
            return None
        raise Exception('illegal type')

    
    def addAdjNodes(self,adjNodes,connectionTypes='bi'):
        if(isinstance(adjNodes,list)):
            if(isinstance(connectionTypes,list)==False):
                connectionTypes =[connectionTypes]

            self.adjNodes = self.adjNodes+adjNodes
            self.adjNodeConnectionTypes = self.adjNodeConnectionTypes + connectionTypes
            for key, adjNode in adjNodes:
                adjNode.adjNodes.append(self)
                if(connectionTypes[key] == 'in'):
                    adjNode.adjNodeConnectionTypes.append('out')
                elif((connectionTypes[key] == 'out')):
                    adjNode.adjNodeConnectionTypes.append('in')
                elif((connectionTypes[key] == 'bi')):
                    adjNode.adjNodeConnectionTypes.append('bi')
                else:
                    raise Exception('adjNodeConnectionTypes should either out, in, or bi.')
        else:
            self.adjNodes.append(adjNodes)
            assert(isinstance(connectionTypes,str))
            self.adjNodeConnectionTypes.append(connectionTypes)
            adjNodes.adjNodes.append(self)
            if(connectionTypes == 'in'):
                adjNodes.adjNodeConnectionTypes.append('out')
            elif((connectionTypes== 'out')):
                adjNodes.adjNodeConnectionTypes.append('in')
            elif((connectionTypes == 'bi')):
                adjNodes.adjNodeConnectionTypes.append('bi')
            else:
                raise Exception('adjNodeConnectionTypes should either out, in, or bi.')


    def delAdjNodes(self,adjNodes):
        #TODO add directional support
        DeprecationWarning('this will be deprecated')
        if(isinstance(adjNodes,list)):
            for adjNode in adjNodes:
                self.adjNodes.remove(adjNode)
                adjNode.adjNodes.remove(self)
        else:
            self.adjNodes.remove(adjNodes)
            adjNodes.adjNodes.remove(self) 
    
    # test whether the adjacent nodes are exclusive
    # TODO tbd
    def testAdjNodesExclusiveness(self):
        for ind, adjnode in self.adjNodes:
            pass
    

    # message-passing-relavent funtions


    def flushMsgs(self):
        '''
            flush all messages
        '''
        self.inMsgs = dict({})
        self.outMsgs = dict({})
    
    def flushConnections(self):
        '''
            flush all connections 
        '''
        self.adjNodeConnectionTypes =[] 
        self.adjNodes = []
        self.flushMsgs()

    def receiveMessage(self,msg,fromNode=None,callback=None,callbackargs=None):
        '''
            receive message from specific port 
            fromNode is index of node in inMsgs
        '''
        assert(isinstance(fromNode,int)==True)
        self.inMsgs[fromNode] = msg
        if(callback!=None):
            callback(callbackargs)

    def sendMsg(self,port=0,callback=None,callbackargs=None):
        '''
            Send message with some specific process
            port: should be an integer
            This PROCESS should be defined in the 
        '''
        self.outMsgs[port] = self.processData(port)
        
        if(callback!=None):
            callback(callbackargs)
            
        return self.outMsgs[port]

    def processData(self,node=None):
        '''
            process received data
            node: 
        '''
        pass
    
    def sendMsgtoNode(self,node,callback=None,callbackargs = None):
        '''
            Send message to specific adjacent node

            node: should be an instance of a node
                 TODO to extended to name and index
        '''
        outMsgs = self.processData(node)



class factNode(cNode):
    def __init__(self,name):
       cNode.__init__(self,name)
       self.type = 'fact'
            

class varNode(cNode):
    def __init__(self,name):
       cNode.__init__(self,name)
       self.type = 'var'       
       self.adjSameChannelInd= [] # store the index of factnode connects the two nodes.

    def getAdjSameNodesIndex(self,adjSameNodes):
        pass
    # TODO should be deprecated
    def addAdjSameNodes(self,varnodes,factIndex =-1):
        #DeprecationWarning('I should deprecate this method')
        if(isinstance(varnodes,list)):
            self.adjSameNodes = self.adjSameNodes+varnodes
            for varnode in varnodes:
                if(varnode.type!= 'var'):
                    raise Exception('Adjacent Same nodes should be a var')
                varnode.adjSameNodes.append(self)
        else:
            if(varnodes.type!= 'var'):
                    raise Exception('Adjacent Same nodes should be a var')
            self.adjSameNodes.append(varnodes)
            varnodes.adjSameNodes.append(self)
    
    def delAdjSameNodes(self,varnodes):
        if(isinstance(varnodes,list)):
            for varnode in varnodes:
                self.adjSameNodes.remove(varnode)
                adjNode.adjSameNodes.remove(self)
        else:
            self.adjSameNodes.remove(varnodes)
            adjNodes.adjSameNodes.remove(self)
    
    #add adjSameNodes with a fact node between it 
    #If the coorresponding varnode is itself, then the factnodes will be a single-port node
    #an ajacent varnode corresponds a single fact while the factnodes can connnect several varnodes
    def addAdjSameNodesWithFact(self,varnodes,factnodes):
        if(isinstance(factnodes,list)==False):
            factnodes = [factnodes]
        for ind,factnode in factnodes:
            self.addAdjNodes(factnode)
            factIndex = len(self.addAdjNodes)-1
            self.addAdjSameNodes(varnodes[ind],factIndex)
            if(isinstance(varnodes[ind],list)==False):
                tmpvarnodes =  [varnodes[ind]]
            else:
                tmpvarnodes = varnodes[ind]

            for varnode in tmpvarnodes:
                if(varnode != self):
                    varnode.addAdjNodes(factnode)

            factnode.addAdjNodes(self)
            factnode.addAdjNodes(tmpvarnodes)
    
# some mission-spec node class
class agentVarNode(varNode):
    '''
        Agent node class  in HSPAWN, each node represents the present state of an agent.
        
        Agent node should recompute its belief when it intends to transmit it. However, when to do it is a question.
        The blowing is severval candicates:
        1.once it sends msgs to. But I should find whether the belief should be updated.
            when one receives new messages, its belief become 'old'
        2.once it receives enough messages
        beliefs: a normalDistributeMsg

        Selection 1 is fine, I should do it first        
    '''
    def __init__(self,name):
        varNode.__init__(self,name)
        self.beliefs = None
        self.beliefisnew = False

    def afterReceivedMsg(self,paras):
        '''
            function called after receiving message
            paras : dict of parameters 

            How to recompute the belief?
            belief should be recomputed when
            1.receive message from mobilityFactNode
            2.have received all p2p messages

            This model has some practical problem yet. However, it is suitable for theoretical analysis.
        '''
        self.beliefisnew  = False

        # if(isinstance(paras['fromnode'],mobilityFactNode)):
        #     logging.debug('I get my beliefs--'+self.name)
        #     self.beliefs = paras['msg']
        # elif len(self.adjNodes) == len(self.inMsgs):
        #     logging.debug( "I should update my beliefs"+self.name)

        #     #self.recomputeBeliefs()
        #     #self.flushMutualMsgs()
        # else:
        #     # do nothiBeliefs
        #     #print 'satllite Here'
        #     pass

    def getBeliefs(self,port = None):
        if(self.beliefisnew == False or port is not None):
            if self.beliefisnew == False and port is None:
                logging.debug("------>\n %s says: I need to update my belief."%self.name)
                logging.debug('previous beliefs')
                if self.beliefs is not None:
                    logging.debug(self.beliefs.mean)
                    logging.debug(self.beliefs.cov)
                else:
                    logging.debug('None')

            msgs = []
            for key in self.inMsgs:
                if port is not None and key == port:
                    continue
                msgs.append(self.inMsgs[key])
                logging.debug(self.inMsgs[key].__str__())
            assert len(msgs) != 0 

            #BUG: in backfarward mode, msgs may contain only chimsg
            msgs.sort()
            if isinstance(msgs[0],normalDistributeMsg) == False :
                logging.debug("I just meet no Gaussian msg problem")
                if port is not None:# have no mormalized msgs
                    partialBelifs = self.getBeliefs()
                    if partialBelifs.dim() in (3,6):
                        tempMsg = msgs[0].approximation(partialBelifs.mean[0:3])
                        if partialBelifs.dim()==6:
                            tempMean = np.zeros((6,))
                            tempMean[0:3] = tempMsg.mean
                            tempMean[3:6] = partialBelifs.mean[3:6]
                            tempCov  = np.zeros((6,6))
                            tempCov[0:3,0:3] = tempMsg.cov
                            tempCov[3:6,3:6] = partialBelifs.cov[3:6,3:6] #NOTE ignoring correlation
                            tempMsg.mean = tempMean
                            tempMsg.cov  = tempCov
                        msgs[0] = tempMsg
                    else:
                        raise Exception("msgs' dimension is not supported yet")
                else:
                    raise Exception("msgs are all Non-Gaussian")
            
            beliefs = mergeMsgs(msgs,True)
            if port is None: 
                self.beliefs = beliefs
                self.beliefisnew = True
                logging.debug('current beliefs: ')
                logging.debug(self.beliefs.mean)
                logging.debug(self.beliefs.cov)
                logging.debug('<------------')
        else:
            beliefs = self.beliefs
        if port is not None:
            logging.debug(self.name + " outmsg is:"+ beliefs.__str__())
        return beliefs
        
    def processData(self,port=None):
        '''
            transmit message to adjcantNode
            TODO some times we should  determite transmiting data type according to receiving node
            As the strategy says, we should just recompute the belief
        '''
        assert  port is not None ,'dont support this argument'
        return self.getBeliefs(port)

    # message-passing-relavent funtions
    def receiveMessage(self,msg,fromNode=None):
        paras = {}
        paras['fromnode'] = self.adjNodes[fromNode]
        paras['msg']      = msg
        
        cNode.receiveMessage(self,msg,fromNode,self.afterReceivedMsg,paras)
    
    def flushMutualMsgs(self):
        '''
            flush all the Mutual msgs for next iteration
        '''
        logging.debug( 'flushing inMsgs..  --- '+self.name)
        index = 0
        for node in self.adjNodes:
            if isinstance(node,mutualConNode):
                self.inMsgs.pop(index)
            index += 1
        logging.debug( 'flushed,len of inMsg:'+ str(len(self.inMsgs)))
        logging.debug( 'len of adjnodes:'+   str(len(self.adjNodes)))

class priorFactNode(factNode):
    '''
        var node indicates the priori distribution
        It just send messages to the outMsgs
    '''
    def setPriorMsg(self,msg):
        self.inMsgs[0] = msg

    def processData(self,node=None):
        return self.inMsgs[0]


class mobilityFactNode(factNode):
    '''
        Fact node represents the mobility. In “Hybrid cooperative positioning based on distributed belief propagation,” 
        Fig 1. it is 'f' node.
        According to Eq(14), msg it sends is a type of normal distribution. the state is a D+1 dimension variable, of which first three
        elements are position and the last one is clock drift
    '''
    def __init__(self,name):
        factNode.__init__(self,name)

    def setMobilityModel(self,model):
        '''
            Set a mobility model for this node.
            model: a function handle. It should be a map from the previous state to the present state.
                   SO, THIS IS A PARTICLE-FORM COMPUTATION
        '''
        self.model = model

    def setVaritionalMobilityModel(self,model):
        '''
            Set a varitional mobility model for this node.
            model is an instance of  
        '''
        self.varitionalModel = model

    def processData(self,port=None):
        '''
            if port isn't None, then process data will process data for specific port
        '''
        # fetch received msg

        #only one in msg
        if self.model.isforward:
            msgIn  = self.inMsgs[0]
        else:
            msgIn  = self.inMsgs[1]

        msgOut = self.model.transform(msgIn)
    
        assert(isinstance(msgIn,normalDistributeMsg))

        return msgOut

class satelliteFactNode(factNode):
    '''
        fact node indicates the priori distribution
        It just send messages to the outMsgs
    '''
    def setPriorMsg(self,msg):
        self.inMsgs[0] = msg

    def processData(self,node=None):
        return self.inMsgs[0]

class mutualConNode(factNode):
    '''
        fact node between agents, this class supports bidirectional schedule
        That means it won't be disabled untill finishing two directional transporting

        measurement  =  link's distance measurement 
        measurementVar = variance of measurement    
    '''

    def __init__(self,name,measurement= None,measurementVar=None):
        factNode.__init__(self,name)
        self.measurement   = measurement
        self.measurementVar= measurementVar
        

    def setMeasuremnt(self,measurement,measurementVar):
        self.measurement   = measurement
        self.measurementVar= measurementVar

    def processData(self,port=None):
        logging.debug( 'I am going to send data ---'+self.name)
        tempSendMsg = self.inMsgs[1-port]
        logging.debug( 'node to: '+ self.adjNodes[port].name)
        mean        = tempSendMsg.mean
        cov  = tempSendMsg.cov + self.measurementVar*np.eye(tempSendMsg.cov.shape[0])

        outMsgName  = self.name+'to'+self.adjNodes[port].name
        varrho      = self.measurement
        tempMsg     = chiMsg(outMsgName,mean,cov,varrho)
        return  tempMsg

    def generateMeasurement(self):
        '''
            generateMeasurement according to parameters
        '''
        pass

class mutualConWithDriftNode(mutualConNode):
    def setFirstLink(self,isFirstLink=None):
        '''
            set as first link
        '''
        if(isFirstLink is None):
            isFirstLink= True
        self.isFirstLink  =  isFirstLink

    def processData(self,port = None):
        '''
            process data to generate different messages w.r.t different port
        '''
        assert port is not None,'dont support this argument'

        # determine receiving node type 
        if isinstance(self.adjNodes[port],timediffVarNode):
            # tranamit to timeDiff
            logging.debug(self.name+' says: I should pass messages to timeDiff---->')
            mobileMsg  =[]
            for key, msg  in self.inMsgs.iteritems():
                if key == port:
                    continue
                if isinstance(self.adjNodes[key],agentVarNode) and isinstance(self.adjNodes[key],timediffVarNode) ==False:
                    mobileMsg.append(msg)
            assert(len(mobileMsg) ==2)
            mobileMsg1  = mobileMsg[0]
            mobileMsg2  = mobileMsg[1]
            logging.debug('the two agents\' position:')
            logging.debug(mobileMsg1.mean[0:3])
            logging.debug(mobileMsg2.mean[0:3])

            u           =  mobileMsg2.mean[0:3]-mobileMsg1.mean[0:3]
            mean        = np.linalg.norm(u)
            logging.debug('the distance is :')
            logging.debug(mean)
            if abs(mean-44078856.0)<1:
                print "b"
            logging.debug('the measurment is:')
            logging.debug(self.measurement)

            covdelta         = mobileMsg1.cov[0:3,0:3]  + mobileMsg2.cov[0:3,0:3]
            sigma2      = 0 
            for  i  in range(len(u)):
                for  j  in range(len(u)):
                    sigma2  += u[i]*u[j]*covdelta[i,j]
            sigma2  = sigma2 / (mean**2)
            if(self.isFirstLink):
                mean    =  mean - self.measurement
            else:
                mean    =  self.measurement - mean
            
            # whether sigma2 should add noise of the measurement?
            sigma2  = sigma2 + self.measurementVar
            tempMsg = normalDistributeMsg('nmsg',mean,sigma2)
            logging.debug('returned message is:')
            logging.debug(tempMsg.__str__())

        elif isinstance(self.adjNodes[port],agentVarNode):
            # get other mobilityFactNode's msg
            logging.debug(self.name+' says: I should pass messages to PosteriorNode----->')
            for key, msg  in self.inMsgs.iteritems():
                if key == port:
                    continue
                if isinstance(self.adjNodes[key],agentVarNode) and isinstance(self.adjNodes[key],       timediffVarNode) ==False:
                    mobileMsg = msg
                if isinstance(self.adjNodes[key],timediffVarNode):
                    timeDiffMsg = msg
            assert((mobileMsg is None)==False)
            assert((timeDiffMsg is None)==False)

            mean        = mobileMsg.mean[0:3]
            
            if(self.isFirstLink):
                varrho      = self.measurement+ timeDiffMsg.mean
            else:
                varrho      = self.measurement- timeDiffMsg.mean

            tempcov     = mobileMsg.cov[0:3,0:3]
            covshape    = tempcov.shape
            cov         = tempcov+ np.eye(covshape[0])*(self.measurementVar+ timeDiffMsg.cov)
            tempMsg     = chiMsg('chimsg',mean,cov,varrho)
            logging.debug('returned message is:')
            logging.debug(tempMsg.__str__())
        else:
            raise Exception('something going wrong')

        logging.debug(self.name+'<<<<<-----')
        return tempMsg

class timediffVarNode(agentVarNode):
    '''
        Right now, it is copy of the agentVarNode
    '''
    pass

